package com.gy.interceptor;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@Component
@Slf4j
public class RateLimitInterceptor implements HandlerInterceptor {

    // 全局限流：全局令牌桶（例如每秒允许 70 次请求，根据服务器性能调整）
    private static final RateLimiter GLOBAL_LIMITER = RateLimiter.create(70);

    // IP 限流（滑动窗口，1秒内最多10次）
    private final Cache<String, AtomicInteger> ipRequestCounts = CacheBuilder.newBuilder()
            .expireAfterWrite(1, TimeUnit.SECONDS)
            .build();
    private static final int MAX_IP_REQUESTS_PER_SECOND = 10;

    @Override
    public boolean preHandle(jakarta.servlet.http.HttpServletRequest request, jakarta.servlet.http.HttpServletResponse response, Object handler) throws Exception {
        String uri = request.getRequestURI();
        // 只对分片上传接口（假设路径以 "/upload/chunk" 开头）应用全局和 IP 限流
        if (uri.startsWith("/upload/chunk")) {
            // 全局限流
            if (!GLOBAL_LIMITER.tryAcquire()) {
                log.warn("全局限流触发，服务器繁忙，URI: {}", uri);
                sendError((HttpServletResponse) response, "请求过于频繁，请稍后再试（全局限流）");
                return false;
            }
            // IP 限流（滑动窗口）
            String ip = getClientIp((HttpServletRequest) request);
            AtomicInteger ipCounter = ipRequestCounts.get(ip, () -> new AtomicInteger(0));
            if (ipCounter.incrementAndGet() > MAX_IP_REQUESTS_PER_SECOND) {
                log.warn("IP [{}] 访问过于频繁，已被限流，URI: {}", ip, uri);
                sendError((HttpServletResponse) response, "请求过于频繁，请稍后再试（IP限流）");
                return false;
            }
        }
        return true;
    }


    /**
     * 获取客户端 IP（防止伪造）
     */
    private String getClientIp(HttpServletRequest request) {
        String[] headers = {"X-Forwarded-For", "X-Real-IP", "Proxy-Client-IP", "WL-Proxy-Client-IP"};
        for (String header : headers) {
            String ip = request.getHeader(header);
            if (ip != null && !ip.isEmpty() && !"unknown".equalsIgnoreCase(ip)) {
                return ip.split(",")[0].trim();
            }
        }
        return request.getRemoteAddr();
    }

    /**
     * 统一返回错误提示给前端，格式为 JSON
     */
    private void sendError(HttpServletResponse response, String message) throws Exception {
        response.setStatus(429);
        response.setContentType("application/json;charset=UTF-8");
        PrintWriter writer = response.getWriter();
        writer.write("{\"error\":\"" + message + "\"}");
        writer.flush();
    }
}
